import numpy as np
import matplotlib.pyplot as plt
import numpy.random as npr
import matplotlib

matplotlib.use(backend="TkAgg")


def experiment():
    '''
    Suppose we draw a random 5-card hand from
    a standard 52-card deck. We want to find the distribu-
    tion of the number of aces in the hand. Let X = #aces.
    We want to determine the PMF of X (or the CDF—but
    the PMF is easier). We know that P (X = k) = 0 ex-
    cept if k = 0, 1, 2, 3, 4. This is clearly not binomial since
    the trials (of drawing cards) are not independen
    '''

    N, K, n = 52, 4, 5
    trials = 100_000

    # 模拟数据：从 N 张牌中不放回抽 n 张，统计 A 的数量
    samples_hyper = [
        np.sum(np.random.choice(a=[1] * K + [0] * (N - K), size=n, replace=False))
        for _ in range(trials)
    ]

    npr.hypergeometric()

    samples_hyper = np.array(samples_hyper)

    theoretical_mean_hyper = n * K / N

    empirical_mean_hyper = sum(samples_hyper) / trials

    # 画图
    plt.figure(figsize=(6, 4))
    plt.hist(samples_hyper, bins=range(n + 2), density=True, alpha=0.6, color='lightcoral', edgecolor='black')
    plt.axvline(theoretical_mean_hyper, color='red', linestyle='--',
                label=f'Theoretical E[X]={theoretical_mean_hyper:.2f}')
    plt.axvline(empirical_mean_hyper, color='green', linestyle=':', label=f'Empirical mean={empirical_mean_hyper:.2f}')
    plt.title(f'Hypergeometric Distribution (N={N}, K={K}, n={n})')
    plt.xlabel("X")
    plt.ylabel("Frequency")
    plt.legend()
    plt.show()


# number of good, number of bad, and number of samples
def test01():
    s = npr.hypergeometric(ngood=10, nbad=5, nsample=10, size=10000)
    # print(s)
    plt.hist(x=s)
    plt.show()


'''
Suppose you have an urn with 15 white and 15 black marbles.
If you pull 15 marbles at random, how likely is it that
12 or more of them are one color?
'''


def test02():
    s = npr.hypergeometric(ngood=15, nbad=15, nsample=15, size=100000)
    ans = sum(s >= 12) / 100000 + sum(s <= 3) / 100000

    print(ans)


if __name__ == '__main__':
    # test01()
    test02()